import openai
import requests
import json
import os
import re
from typing import List, Dict, Any, Optional

# API设置
OPENAI_API_KEY = "your-api-key"  # 替换为你的OpenAI API密钥
OPENAI_API_BASE = "https://www.dmxapi.com/v1"  # 可以替换为其他API端点
SERPER_API_KEY = "your-serper-key"  # 替换为你的Serper API密钥
SERPER_API_URL = "https://google.serper.dev/search"

# 配置OpenAI
openai.api_key = OPENAI_API_KEY
if OPENAI_API_BASE:
    # 检查OpenAI库版本
    if hasattr(openai, "base_url"):  # v1.0.0+
        openai.base_url = OPENAI_API_BASE
    else:  # v0.x
        openai.api_base = OPENAI_API_BASE

class LinkFinder:
    def __init__(self, openai_api_key: str, serper_api_key: str):
        """初始化链接查找器"""
        self.openai_api_key = openai_api_key
        self.serper_api_key = serper_api_key
        
    def search_google(self, query: str, num_results: int = 5) -> Dict[str, Any]:
        """使用Serper API搜索Google"""
        headers = {
            'X-API-KEY': self.serper_api_key,
            'Content-Type': 'application/json'
        }
        
        payload = {
            'q': query,
            'num': num_results
        }
        
        try:
            response = requests.post(
                SERPER_API_URL, 
                headers=headers,
                json=payload
            )
            response.raise_for_status()  # 检查HTTP错误
            return response.json()
        except Exception as e:
            print(f"搜索错误: {e}")
            return {"error": str(e)}
    
    def extract_links(self, search_results: Dict[str, Any]) -> List[Dict[str, str]]:
        """从搜索结果中提取链接"""
        links = []
        
        # 处理搜索结果中的organic部分
        if "organic" in search_results:
            for item in search_results["organic"]:
                if "title" in item and "link" in item:
                    link_data = {
                        "title": item["title"],
                        "link": item["link"],
                        "snippet": item.get("snippet", "")
                    }
                    links.append(link_data)
        
        # 处理可能的知识面板结果
        if "knowledgeGraph" in search_results:
            kg = search_results["knowledgeGraph"]
            if "title" in kg and "url" in kg:
                link_data = {
                    "title": f"知识图谱: {kg['title']}",
                    "link": kg["url"],
                    "snippet": kg.get("description", "")
                }
                links.append(link_data)
        
        return links
    
    def generate_search_queries(self, text: str, num_queries: int = 1) -> List[str]:
        """使用GPT生成搜索查询"""
        try:
            # 构建系统提示和用户提示
            system_prompt = """你是一个专业的搜索查询生成器。
            根据给定的文本，生成最佳的Google搜索查询，以便找到与文本内容相关的可靠来源链接。
            返回的查询应该是高质量的，能够找到权威的、可信的信息来源。
            只返回搜索查询本身，不要包含任何其他文字或解释。"""
            
            user_prompt = f"""请为以下文本生成{num_queries}个最佳搜索查询，每行一个查询：
            
            文本：{text}
            
            查询："""
            
            # 检测OpenAI API版本并调用
            if hasattr(openai, "chat") and hasattr(openai.chat, "completions"):
                # 新版本API (v1.0.0+)
                response = openai.chat.completions.create(
                    model="gpt-4o",
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt}
                    ],
                    temperature=0.7,
                    max_tokens=100
                )
                queries_text = response.choices[0].message.content.strip()
            else:
                # 旧版本API (v0.x)
                response = openai.ChatCompletion.create(
                    model="gpt-4o",
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt}
                    ],
                    temperature=0.7,
                    max_tokens=100
                )
                queries_text = response["choices"][0]["message"]["content"].strip()
            
            # 提取并处理查询
            queries = [q.strip() for q in queries_text.split('\n') if q.strip()]
            
            return queries[:num_queries]  # 确保不超过请求的查询数量
            
        except Exception as e:
            print(f"生成查询时出错: {e}")
            # 如果API调用失败，使用简单方法生成查询
            return [text]
    
    def analyze_and_rank_links(self, text: str, links: List[Dict[str, str]]) -> List[Dict[str, str]]:
        """使用GPT分析和排名链接的相关性"""
        if not links:
            return []
            
        try:
            # 准备链接数据用于分析
            links_data = "\n".join([
                f"{i+1}. 标题: {link['title']}\n   链接: {link['link']}\n   摘要: {link['snippet']}"
                for i, link in enumerate(links)
            ])
            
            # 构建系统提示
            system_prompt = """你是一个专业的信息源评估专家。
            你的任务是评估提供的链接与给定文本的相关性和可靠性。
            对于每个链接，考虑以下因素：
            1. 相关性 - 链接内容与文本的相关程度
            2. 可靠性 - 链接源的权威性和可信度
            3. 时效性 - 信息是否是最新的（如果适用）
            请排序这些链接，将最相关、最可靠的链接排在前面。"""
            
            # 构建用户提示
            user_prompt = f"""请评估以下链接与文本的相关性，并按照相关性从高到低排序：
            
            文本：{text}
            
            链接列表：
            {links_data}
            
            请返回排序后的链接索引列表，格式为：[3, 1, 5, 2, 4]（表示第3个链接最相关，第1个次之，依此类推）。
            只返回索引列表，不要包含任何其他文字。"""
            
            # 检测OpenAI API版本并调用
            if hasattr(openai, "chat") and hasattr(openai.chat, "completions"):
                # 新版本API (v1.0.0+)
                response = openai.chat.completions.create(
                    model="gpt-4o",
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt}
                    ],
                    temperature=0.3,
                    max_tokens=50
                )
                result = response.choices[0].message.content.strip()
            else:
                # 旧版本API (v0.x)
                response = openai.ChatCompletion.create(
                    model="gpt-4o",
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt}
                    ],
                    temperature=0.3,
                    max_tokens=50
                )
                result = response["choices"][0]["message"]["content"].strip()
            
            # 解析返回的索引列表
            try:
                # 尝试解析为JSON数组
                indices = json.loads(result)
            except json.JSONDecodeError:
                # 如果不是有效的JSON，尝试使用正则表达式
                indices_match = re.search(r'\[([\d, ]+)\]', result)
                if indices_match:
                    indices_str = indices_match.group(1)
                    indices = [int(idx.strip()) for idx in indices_str.split(',') if idx.strip().isdigit()]
                else:
                    # 如果仍然失败，默认使用原始顺序
                    return links
            
            # 验证索引的有效性并调整为0基索引
            valid_indices = []
            for idx in indices:
                idx_zero_based = idx - 1  # 转换为0基索引
                if 0 <= idx_zero_based < len(links):
                    valid_indices.append(idx_zero_based)
            
            # 根据排序结果重新排列链接
            sorted_links = []
            for idx in valid_indices:
                sorted_links.append(links[idx])
            
            # 添加任何未包含在结果中的链接
            included_indices = set(valid_indices)
            for i in range(len(links)):
                if i not in included_indices:
                    sorted_links.append(links[i])
            
            return sorted_links
            
        except Exception as e:
            print(f"分析链接时出错: {e}")
            return links  # 发生错误时返回原始链接列表
    
    def find_relevant_links(self, text: str, num_results: int = 5) -> List[Dict[str, str]]:
        """查找与文本相关的链接"""
        # 生成搜索查询
        queries = self.generate_search_queries(text, num_queries=2)
        
        all_links = []
        # 对每个查询执行搜索
        for query in queries:
            # 搜索Google
            search_results = self.search_google(query, num_results=num_results)
            
            # 提取链接
            links = self.extract_links(search_results)
            all_links.extend(links)
        
        # 去除重复链接
        unique_links = []
        seen_urls = set()
        for link in all_links:
            if link["link"] not in seen_urls:
                seen_urls.add(link["link"])
                unique_links.append(link)
        
        # 分析和排序链接
        if len(unique_links) > num_results:
            ranked_links = self.analyze_and_rank_links(text, unique_links)
            return ranked_links[:num_results]
        else:
            return unique_links

# 使用示例
def main():
    # 使用环境变量或直接设置API密钥
    openai_api_key = os.environ.get("OPENAI_API_KEY", OPENAI_API_KEY)
    serper_api_key = os.environ.get("SERPER_API_KEY", SERPER_API_KEY)
    
    # 创建链接查找器
    finder = LinkFinder(openai_api_key, serper_api_key)
    
    # 示例文本
    text = input("请输入要查找链接的文本: ")
    
    # 查找相关链接
    links = finder.find_relevant_links(text, num_results=3)
    
    # 打印结果
    print("\n找到的相关链接:")
    for i, link in enumerate(links, 1):
        print(f"\n{i}. {link['title']}")
        print(f"   URL: {link['link']}")
        print(f"   摘要: {link['snippet']}")

if __name__ == "__main__":
    main()
